package io.sqrtqiezi.spark.streaming

import org.apache.spark.SparkConf
import org.apache.spark.streaming.{Seconds, StreamingContext}

object StateTracker1 {
  def main(args: Array[String]): Unit = {
    val conf: SparkConf = new SparkConf()
      .setAppName("spark state tracker1")
      .setMaster("local[2]")

    val ssc = new StreamingContext(conf, Seconds(5))
    ssc.checkpoint("lagou-data/checkpoint1/")

    val wordDStream = ssc.socketTextStream("localhost", 9999)
      .flatMap(_.split("\\s+"))
      .map((_, 1))

    val updateFunc = (currentValues: Seq[Int], prevValueState: Option[Int]) => {
      val currentCount = currentValues.sum
      val previousCount = prevValueState.getOrElse(0)
      Some(currentCount + previousCount)
    }

    val stateDStream = wordDStream.updateStateByKey[Int](updateFunc)
    stateDStream.print()

    stateDStream.repartition(1)
      .saveAsTextFiles("data/output1")

    ssc.start()
    ssc.awaitTermination()
  }
}
